`timescale 1ns / 1ps
`include "define.sv" // Ensure this file defines `m32n8k16`, `INT4`, etc.

// Simplified ram_32x1024_burst module for testbench simulation
// For actual simulation, please use your real ram_32x1024_burst.sv file
/* module ram_32x1024_burst #(
    parameter DATA_WIDTH = 32,
    parameter ADDR_WIDTH = 10
)(
    input wire clk,
    input wire rst_n,
    input wire we,      // Write Enable
    input wire re,      // Read Enable (added)
    input wire burst_en,
    input wire [ADDR_WIDTH-1:0] addr,
    input wire [DATA_WIDTH-1:0] din,
    input wire [7:0] burst_len,
    output reg [DATA_WIDTH-1:0] dout,
    output reg burst_done
);
    reg [DATA_WIDTH-1:0] mem [0:(1<<ADDR_WIDTH)-1]; // Internal memory

    reg [7:0] write_counter;
    reg [7:0] read_counter; // Read counter (added)
    reg write_active;
    reg read_active;    // Read active signal (added)

    always @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            write_counter <= 0;
            read_counter <= 0;
            write_active <= 0;
            read_active <= 0;
            burst_done <= 0;
        end else begin
            // Write Logic
            if (we && burst_en && !write_active) begin // Burst write start
                write_active <= 1;
                write_counter <= 0;
                burst_done <= 0;
            end else if (write_active) begin
                if (write_counter <= burst_len) begin
                    mem[addr + write_counter] <= din;
                    write_counter <= write_counter + 1;
                    burst_done <= 0;
                end else begin
                    write_active <= 0;
                    burst_done <= 1; // Write burst complete
                end
            end else begin
                burst_done <= 0; // Default de-assert
            end

            // Read Logic (added)
            if (re && burst_en && !read_active) begin // Burst read start
                read_active <= 1;
                read_counter <= 0;
            end else if (read_active) begin
                if (read_counter <= burst_len) begin
                    dout <= mem[addr + read_counter];
                    read_counter <= read_counter + 1;
                end else begin
                    read_active <= 0;
                    dout <= '0; // Clear output after read burst
                end
            end else if (!write_active) begin // Ensure dout is stable when not writing
                dout <= mem[addr]; // Simulate single cycle read or hold old value
            end
        end
    end
endmodule */


module tb_memory_controller_axi_full;

    // Parameters for memory_controller instance
    localparam DATA_WIDTH = 32;
    localparam ADDR_WIDTH = 10; // Address width of the RAM module

    // AXI4-Full Master Signals
    // AXI Global Signals
    reg ACLK;
    reg ARESETN;

    // AW Channel (Address Write)
    reg [3:0] aw_id;
    reg [31:0] aw_addr;
    reg [7:0] aw_len;
    reg aw_valid;
    wire aw_ready;

    // W Channel (Write Data)
    reg [3:0] w_id;
    reg [DATA_WIDTH-1:0] w_data;
    reg w_last;
    reg w_valid;
    wire w_ready;

    // B Channel (Write Response)
    wire b_valid;
    reg b_ready;

    // AR Channel (Address Read)
    reg [3:0] ar_id;
    reg [31:0] ar_addr;
    reg [7:0] ar_len;
    reg ar_valid;
    wire ar_ready;

    // R Channel (Read Data)
    wire [3:0] r_id;
    wire [DATA_WIDTH-1:0] r_data;
    wire r_last;
    wire r_valid;
    reg r_ready;

    // Internal signals connected to memory_controller (derived from AXI signals)
    wire internal_we = w_valid & w_ready; // Write enable
    wire internal_re; // Read enable
    wire [ADDR_WIDTH-1:0] internal_addr = ar_valid ? ar_addr[ADDR_WIDTH-1:0] : aw_addr[ADDR_WIDTH-1:0]; // Mux read/write address
    wire [DATA_WIDTH-1:0] internal_din = w_data;
    wire [7:0] internal_burst_len = ar_valid ? ar_len : aw_len; // Mux read/write burst length

    // Memory controller inputs
    reg [1:0] matrix_type;
    reg [1:0] data_type;

    // Memory controller outputs - 修改为多个RAM输出
    wire [DATA_WIDTH-1:0] ram0_dout, ram1_dout, ram2_dout;
    wire ram0_burst_done, ram1_burst_done, ram2_burst_done;
    
    // 选择当前读取的RAM数据输出
    wire [DATA_WIDTH-1:0] ram_dout;
    wire burst_done;
    
    assign ram_dout = (u_ram_ctrl.read_ram == 2'b00) ? ram0_dout :
                      (u_ram_ctrl.read_ram == 2'b01) ? ram1_dout :
                      (u_ram_ctrl.read_ram == 2'b10) ? ram2_dout : 32'h0;
                      
    assign burst_done = (u_ram_ctrl.read_ram == 2'b00) ? ram0_burst_done :
                        (u_ram_ctrl.read_ram == 2'b01) ? ram1_burst_done :
                        (u_ram_ctrl.read_ram == 2'b10) ? ram2_burst_done : 1'b0;

    // AXI Handshake (Testbench as Master, Slave is always ready)
    assign aw_ready = 1'b1; // Simplified: Slave AW channel is always ready
    assign w_ready = 1'b1;  // Simplified: Slave W channel is always ready
    assign b_valid = aw_valid && w_valid && w_last; // Simplified: Write response valid after last data beat
    assign ar_ready = 1'b1; // Simplified: Slave AR channel is always ready
    assign r_id = ar_id;    // Simplified: Read response ID

    // **修改：r_last 和 r_valid 基于新的内部信号**
    assign r_last = u_ram_ctrl.read_active && (u_ram_ctrl.current_read_beat == internal_burst_len); // r_last based on internal read length
    assign r_valid = u_ram_ctrl.read_active; // Read data valid when internal read is active

    assign internal_re = ar_valid & ar_ready; // Read enable trigger

    // Golden Reference Memory Model
    // Assuming ADDR_WIDTH max is 10, corresponding to 1024 addresses
    reg [DATA_WIDTH-1:0] golden_mem [0:(1<<ADDR_WIDTH)-1];

    // Instantiate the memory_controller - 修改端口连接
    memory_controller #(
        .DATA_WIDTH(DATA_WIDTH),
        .ADDR_WIDTH(ADDR_WIDTH)
    ) u_ram_ctrl (
        .clk(ACLK),
        .rst_n(ARESETN),
        .matrix_type(matrix_type),
        .data_type(data_type),
        .we(internal_we),
        .re(internal_re),
        .addr(internal_addr),
        .din(internal_din),
        .burst_len(internal_burst_len),
        .ram0_dout(ram0_dout),
        .ram1_dout(ram1_dout),
        .ram2_dout(ram2_dout),
        .ram0_burst_done(ram0_burst_done),
        .ram1_burst_done(ram1_burst_done),
        .ram2_burst_done(ram2_burst_done)
     );

    // Clock generation
    always #5 ACLK = !ACLK; // 100 MHz clock (10 ns period)

    // AXI Write Burst Task
    task axi_write_burst;
        input [31:0] address;
        input [7:0] length; // burst_len (0-255, means length+1 transfers)
        input [1:0] mat_type;
        input [1:0] d_type;
        input [DATA_WIDTH-1:0] start_data;
        reg [DATA_WIDTH-1:0] current_data;
    begin
        matrix_type = mat_type;
        data_type = d_type;

        $display("----------------------------------------------------------------------------------");
        $display("Time: %0t, Starting AXI Write Burst. Addr: 0x%H, Len: %0d, Matrix Type: %0b, Data Type: %0b",
                 $time, address, length, mat_type, d_type);

        // AW Phase
        aw_id = 4'b0;
        aw_addr = address;
        aw_len = length;
        aw_valid = 1;

        // W Phase Initialization
        w_id = 4'b0;
        w_valid = 0; // 初始设置为0
        b_ready = 1; // Ready to receive write response

        current_data = start_data;

        // Wait for AXI AW Handshake to complete
        @(posedge ACLK);
        while (~aw_ready) @(posedge ACLK); // Wait for Slave AW Ready

        // W Channel Transfers
        for (int i = 0; i <= length; i = i + 1) begin
            w_data = current_data;
            w_last = (i == length); // Last transfer
            w_valid = 1; // 在循环中设置为1

            @(posedge ACLK);
            // Wait for AXI W Handshake to complete
            while (~w_ready) begin
                @(posedge ACLK);
            end

            // Update Golden Reference Memory Model
            golden_mem[address[ADDR_WIDTH-1:0] + i] = w_data;

            $display("Time: %0t, AXI Write: W_DATA=0x%H, W_LAST=%0b, RAM_WE=%0b, current_ram=%0b, bit_counter=%0d",
                     $time, w_data, w_last, internal_we, u_ram_ctrl.current_ram, u_ram_ctrl.bit_counter);
            current_data = current_data + 1; // Simulate data increment
        end

        // End AXI Transfer
        @(posedge ACLK);
        aw_valid = 0;
        w_valid = 0;
        w_last = 0; // Reset w_last
        // Wait for write response
        while(~b_valid) @(posedge ACLK);
        @(posedge ACLK); // Receive response
        b_ready = 0; // Complete write response reception

        $display("Time: %0t, AXI Write Burst Completed.", $time);
        $display("----------------------------------------------------------------------------------");
    end
    endtask

    // AXI Read Burst Task
    task axi_read_burst;
        input [31:0] address;
        input [7:0] length; // burst_len (0-255, means length+1 transfers)
    begin
        $display("----------------------------------------------------------------------------------");
        $display("Time: %0t, Starting AXI Read Burst. Addr: 0x%H, Len: %0d", $time, address, length);

        // AR Phase
        ar_id = 4'b0;
        ar_addr = address;
        ar_len = length;
        ar_valid = 1;
        r_ready = 1; // Ready to receive read data

        @(posedge ACLK);
        while (~ar_ready) @(posedge ACLK); // Wait for Slave AR Ready

        @(posedge ACLK);  // ← 在这里添加一个周期延迟，让整个循环推迟开始

        // R Channel Transfers and Data Verification
        for (int i = 0; i <= length; i = i + 1) begin
            @(posedge ACLK);
            // Wait for R Handshake
            while (~r_valid) @(posedge ACLK);

            $display("Debug: Reading addr=0x%H, expected=0x%H, actual=0x%H, read_ram=%0d", 
                     address[ADDR_WIDTH-1:0] + i, 
                     golden_mem[address[ADDR_WIDTH-1:0] + i], 
                     ram_dout,
                     u_ram_ctrl.read_ram);

            // Perform Data Verification
            if (ram_dout != golden_mem[address[ADDR_WIDTH-1:0] + i]) begin
                $error("!!! Data Mismatch !!! Time: %0t, Address: 0x%H, Expected Data: 0x%H, Actual Data: 0x%H",
                       $time, address[ADDR_WIDTH-1:0] + i, golden_mem[address[ADDR_WIDTH-1:0] + i], ram_dout);
            end else begin
                $display("Time: %0t, AXI Read: Addr=0x%H, Read Data=0x%H (Verification PASSED)",
                         $time, address[ADDR_WIDTH-1:0] + i, ram_dout);
            end
        end

        // End AXI Read Transfer
        @(posedge ACLK);
        ar_valid = 0;
        r_ready = 0;
        $display("Time: %0t, AXI Read Burst Completed.", $time);
        $display("----------------------------------------------------------------------------------");
    end
    endtask

    // Test Sequence
    initial begin
        ACLK = 0;
        ARESETN = 0;
        // Initialize all AXI signals to 0
        aw_id = 0; aw_addr = 0; aw_len = 0; aw_valid = 0;
        w_id = 0; w_data = 0; w_last = 0; w_valid = 0;
        b_ready = 0;
        ar_id = 0; ar_addr = 0; ar_len = 0; ar_valid = 0;
        r_ready = 0;

        // Initialize Golden Reference Memory to 0
        for (int i = 0; i < (1<<ADDR_WIDTH); i++) begin
            golden_mem[i] = 0;
        end

        #20 ARESETN = 1; // Release reset

        // --- Test Case 1: m32n8k16, INT4 ---
        // switch_threshold = 31 (corresponds to 32 transfers)
        // First burst, should write to RAM0
        axi_write_burst(32'h0000_0000, 8'd31, `m32n8k16, `INT4, 32'h0000_0001); // 32 transfers
        #20; // Wait a few cycles to ensure burst_done stabilizes
        axi_read_burst(32'h0000_0000, 8'd31); // Read back and verify

        #50; // Wait a few cycles

        // Second burst, should write to RAM1 (as bit_counter reached threshold, current_ram incremented)
        axi_write_burst(32'h0000_0000, 8'd31, `m32n8k16, `INT4, 32'h0000_0101); // 32 transfers
        #10;
        axi_read_burst(32'h0000_0000, 8'd31); // Read back and verify

        #50;

        // Third burst, should write to RAM2
        axi_write_burst(32'h0000_0000, 8'd31, `m32n8k16, `INT4, 32'h0000_0201); // 32 transfers
        #10;
        axi_read_burst(32'h0000_0000, 8'd31); // Read back and verify

        #50;

        // End simulation
        $display("--- Simulation Finished ---");
        $finish;
    end

    // Monitor burst_done signal from memory_controller - 修改为监控所有RAM
    always @(posedge ram0_burst_done) begin
        $display("------------------------------------------------------------------");
        $display("!!! RAM0 burst_done asserted by memory_controller at time %0t !!!", $time);
        $display("------------------------------------------------------------------");
    end

    always @(posedge ram1_burst_done) begin
        $display("------------------------------------------------------------------");
        $display("!!! RAM1 burst_done asserted by memory_controller at time %0t !!!", $time);
        $display("------------------------------------------------------------------");
    end

    always @(posedge ram2_burst_done) begin
        $display("------------------------------------------------------------------");
        $display("!!! RAM2 burst_done asserted by memory_controller at time %0t !!!", $time);
        $display("------------------------------------------------------------------");
    end

    // Monitor current_ram and bit_counter
    always @(posedge ACLK) begin
        if (ARESETN) begin
            $display("Time: %0t, current_ram: %0b, read_ram: %0b, bit_counter: %0d, switch_threshold: %0d",
                     $time, u_ram_ctrl.current_ram, u_ram_ctrl.read_ram, u_ram_ctrl.bit_counter, u_ram_ctrl.switch_threshold);
        end
    end

endmodule